from mymodel import * 


shape = (1,2)
x = torch.rand(shape, dtype=torch.float32)
y = torch.rand(shape, dtype=torch.float32)
model = MyModule()
out = torch.onnx.export(
    model,
    (x, y),
    "model.onnx",
    opset_version=15,
    input_names=["x", "y"],
)
# 虽然输出是 ONNX，但内部会先经过 ATen IR 转换。

print(type(out))
print(out)
